#!/usr/bin/env python3

import pyaudio
import json
import time
import os
import requests
import subprocess
import re
import threading
import tempfile
from vosk import Model, KaldiRecognizer

class RoboticArmVoiceControl:
    def __init__(self):
        self.model = Model("../vosk-model-small-cn-0.22")
        self.api_key = "ms-fecb8eca-0503-44e9-9cf2-4870bb0ed37c"
        self.base_url = "https://api-inference.modelscope.cn/v1"
        self.model_name = "deepseek-ai/DeepSeek-V3.2-Exp"
        self.speaking = False
        self.recording = False
        
        # 机械臂状态 (使用0-1000的位置值)
        self.arm_state = {
            "base": 500,      # 1号舵机: 底座 (0-1000)
            "shoulder": 500,  # 2号舵机: 肩部 (0-1000)
            "elbow": 500,     # 3号舵机: 肘部 (0-1000)
            "wrist_vert": 500, # 4号舵机: 腕部垂直 (0-1000)
            "wrist_rot": 500, # 5号舵机: 腕部旋转 (0-1000)
            "gripper": 0      # 6号舵机: 夹爪 (0:闭合, 1000:打开)
        }
        
        # 舵机ID映射
        self.servo_ids = {
            "base": 1,
            "shoulder": 2,
            "elbow": 3,
            "wrist_vert": 4,
            "wrist_rot": 5,
            "gripper": 6
        }
        
        print("🤖 机械臂语音控制系统已启动")
        print("📍 当前舵机状态:", self.arm_state)
        self.speak("系统已启动")
        
    def speak(self, text):
        """改进的语音播报 - 使用更好的语音合成"""
        print(f"🔊 {text}")
        
        if self.speaking:
            return
            
        def _speak():
            self.speaking = True
            try:
                # 方法1: 使用festival（如果可用）
                with open(os.devnull, 'w') as devnull:
                    subprocess.run(["festival", "--tts"], 
                                 input=text.encode('utf-8'),
                                 stdout=devnull, 
                                 stderr=devnull,
                                 timeout=10)
            except FileNotFoundError:
                try:
                    # 方法2: 使用espeak并增加振幅
                    with open(os.devnull, 'w') as devnull:
                        subprocess.run(["espeak", "-v", "zh", "-s", "200", "-a", "200", text], 
                                     stdout=devnull, 
                                     stderr=devnull,
                                     timeout=5)
                except Exception as e:
                    print(f"语音合成错误: {e}")
            except Exception as e:
                print(f"语音合成错误: {e}")
            finally:
                self.speaking = False
        
        # 在新线程中执行语音播报
        thread = threading.Thread(target=_speak)
        thread.daemon = True
        thread.start()
        
    def record_audio_direct(self, duration=5):
        """直接录制音频，不使用静默检测"""
        print(f"🎤 开始录音 {duration}秒...")
        self.speak("请说话")
        
        # 使用arecord直接录制
        try:
            with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_file:
                temp_filename = temp_file.name
            
            # 录制音频 - 使用stdout和stderr重定向替代capture_output
            with open(os.devnull, 'w') as devnull:
                result = subprocess.run([
                    "arecord", 
                    "-f", "S16_LE", 
                    "-r", "16000", 
                    "-c", "1", 
                    "-d", str(duration),
                    temp_filename
                ], stdout=devnull, stderr=devnull, timeout=duration+2)
            
            if result.returncode == 0:
                print("✅ 录音完成")
                return temp_filename
            else:
                print("❌ 录音失败")
                if os.path.exists(temp_filename):
                    os.unlink(temp_filename)
                return None
                
        except subprocess.TimeoutExpired:
            print("❌ 录音超时")
            if os.path.exists(temp_filename):
                os.unlink(temp_filename)
            return None
        except Exception as e:
            print(f"录音错误: {e}")
            if os.path.exists(temp_filename):
                os.unlink(temp_filename)
            return None
    
    def transcribe_audio(self, audio_file):
        """转录音频文件"""
        if not audio_file or not os.path.exists(audio_file):
            return ""
            
        try:
            # 使用vosk进行语音识别
            recognizer = KaldiRecognizer(self.model, 16000)
            
            # 读取音频文件
            with open(audio_file, 'rb') as f:
                audio_data = f.read()
            
            # 分块处理音频
            chunk_size = 4000
            text_result = ""
            
            for i in range(0, len(audio_data), chunk_size):
                chunk = audio_data[i:i+chunk_size]
                if recognizer.AcceptWaveform(chunk):
                    result = recognizer.Result()
                    result_dict = json.loads(result)
                    if result_dict['text']:
                        text_result += result_dict['text'] + " "
            
            # 获取最终结果
            final_result = recognizer.FinalResult()
            final_dict = json.loads(final_result)
            if final_dict['text']:
                text_result += final_dict['text']
            
            # 清理临时文件
            os.unlink(audio_file)
            
            return text_result.strip()
            
        except Exception as e:
            print(f"转录错误: {e}")
            # 清理临时文件
            if os.path.exists(audio_file):
                os.unlink(audio_file)
            return ""
    
    def analyze_voice_command(self, text):
        """分析语音命令并生成舵机控制指令"""
        system_prompt = """你是一个机械臂控制专家。请将用户的语音命令转换为具体的舵机控制指令。

机械臂有6个舵机：
1. base (底座): 0-1000位置，控制整体左右旋转
2. shoulder (肩部): 0-1000位置，控制大臂上下俯仰
3. elbow (肘部): 0-1000位置，控制小臂前后伸展
4. wrist_vert (腕部垂直): 0-1000位置，控制手腕上下摆动
5. wrist_rot (腕部旋转): 0-1000位置，控制手腕水平旋转
6. gripper (夹爪): 0-1000位置，0为闭合，1000为完全打开

请根据用户指令，生成JSON格式的控制指令，包含以下字段：
- action: 动作描述
- commands: 舵机控制列表，每个命令包含 servo(舵机名称) 和 position(目标位置，0-1000)
- duration: 执行时间(秒)

示例指令：
用户说："向右转"
输出：{"action": "底座向右旋转", "commands": [{"servo": "base", "position": 700}], "duration": 2}

用户说："抬起手臂"
输出：{"action": "抬起机械臂", "commands": [{"servo": "shoulder", "position": 300}, {"servo": "elbow", "position": 700}], "duration": 3}

用户说："抓取物体"
输出：{"action": "抓取物体", "commands": [{"servo": "gripper", "position": 0}], "duration": 1}

用户说："旋转手腕"
输出：{"action": "旋转手腕", "commands": [{"servo": "wrist_rot", "position": 800}], "duration": 2}

用户说："放下手腕"
输出：{"action": "放下手腕", "commands": [{"servo": "wrist_vert", "position": 300}], "duration": 2}

请只返回JSON格式，不要其他文字。"""

        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        data = {
            "model": self.model_name,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"当前状态: {self.arm_state}。指令: {text}"}
            ],
            "temperature": 0.3,
            "max_tokens": 500
        }
        
        try:
            print("🔄 解析语音指令...")
            response = requests.post(
                f"{self.base_url}/chat/completions",
                headers=headers,
                json=data,
                timeout=30
            )
            
            if response.status_code == 200:
                result = response.json()
                response_text = result["choices"][0]["message"]["content"].strip()
                
                # 提取JSON部分
                try:
                    command = json.loads(response_text)
                    return command
                except:
                    json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
                    if json_match:
                        command = json.loads(json_match.group())
                        return command
                    else:
                        return {"error": "无法解析AI响应"}
            else:
                return {"error": f"API调用失败: {response.status_code}"}
                
        except Exception as e:
            return {"error": f"API调用异常: {e}"}
    
    def execute_commands(self, command_data):
        """执行舵机控制命令"""
        if "error" in command_data:
            error_msg = f"错误: {command_data['error']}"
            self.speak(error_msg)
            return False
        
        action = command_data.get('action', '未知动作')
        duration = command_data.get('duration', 2)
        
        self.speak(f"执行{action}")
        
        commands = command_data.get("commands", [])
        
        # 生成 movemulti 命令
        positions = []
        for servo_name in self.servo_ids.keys():
            target_pos = None
            for cmd in commands:
                if cmd.get("servo") == servo_name:
                    target_pos = cmd.get("position")
                    break
            
            if target_pos is not None:
                positions.append(str(target_pos))
                self.arm_state[servo_name] = target_pos
            else:
                positions.append("-")
        
        # 发送 movemulti 命令
        movemulti_cmd = f"movemulti {','.join(positions)}"
        print(f"[硬件控制] {movemulti_cmd}")
        self.send_command_to_arm(movemulti_cmd)
        
        # 执行时间
        print(f"   ⏳ 执行中...", end="", flush=True)
        for i in range(duration):
            print(".", end="", flush=True)
            time.sleep(1)
        print(" 完成!")
        
        self.speak("完成")
        print(f"📍 更新后状态: {self.arm_state}")
        return True
    
    def control_single_servo(self, servo_name, position):
        """控制单个舵机"""
        if servo_name in self.servo_ids:
            servo_id = self.servo_ids[servo_name]
            command = f"move {servo_id} {position}"
            print(f"[硬件控制] {command}")
            self.send_command_to_arm(command)
            self.arm_state[servo_name] = position
    
    def send_command_to_arm(self, command):
        """发送命令到机械臂控制器"""
        print(f"[发送命令] {command}")
    
    def voice_menu_selection(self, menu_options, prompt_text):
        """语音菜单选择"""
        while True:
            # 先显示菜单，再语音提示
            print("\n📋 可用选项:")
            for key, description in menu_options.items():
                print(f"  {key}. {description}")
            
            self.speak(prompt_text)
            time.sleep(1)
            
            # 录制音频
            audio_file = self.record_audio_direct(duration=5)
            if audio_file:
                # 转录音频
                choice = self.transcribe_audio(audio_file)
                print(f"📝 识别结果: {choice}")
            else:
                choice = ""
            
            if choice:
                # 中文数字转阿拉伯数字
                chinese_to_arabic = {
                    '一': '1', '二': '2', '三': '3', '四': '4', '五': '5',
                    '六': '6', '七': '7', '八': '8', '九': '9', '十': '10',
                    '1': '1', '2': '2', '3': '3', '4': '4', '5': '5',
                    '6': '6', '7': '7', '8': '8', '9': '9'
                }
                
                # 直接匹配
                clean_choice = choice.replace(' ', '').strip()
                if clean_choice in chinese_to_arabic:
                    choice_num = chinese_to_arabic[clean_choice]
                else:
                    # 提取第一个数字
                    numbers = re.findall(r'\d', clean_choice)
                    if numbers:
                        choice_num = numbers[0]
                    else:
                        # 关键词匹配
                        for key, description in menu_options.items():
                            if key in clean_choice or description in clean_choice:
                                choice_num = key
                                break
                        else:
                            choice_num = None
                
                if choice_num and choice_num in menu_options:
                    selected = menu_options[choice_num]
                    print(f"✅ 选择: {selected}")
                    self.speak(f"选择{selected}")
                    return choice_num
                
                # 检查取消词
                if any(word in choice for word in ["退出", "取消", "返回", "结束"]):
                    return "cancel"
                
                print(f"❌ 无法识别: {choice}")
                self.speak("请重新选择")
            else:
                self.speak("请说话")
    
    def interactive_control(self):
        """交互式控制模式"""
        self.speak("欢迎使用")
        
        while True:
            menu_options = {
                "1": "语音控制",
                "2": "复位机械臂",
                "3": "读取状态",
                "4": "单个舵机控制",
                "5": "退出系统"
            }
            
            print("\n" + "="*40)
            print("🤖 机械臂语音控制系统")
            print("="*40)
            
            choice = self.voice_menu_selection(menu_options, "请选择")
            
            if choice == "1":
                self.voice_control_mode()
            elif choice == "2":
                self.reset_arm()
            elif choice == "3":
                self.read_servo_status()
            elif choice == "4":
                self.single_servo_control()
            elif choice == "5" or choice == "cancel":
                self.speak("再见")
                break
            else:
                self.speak("请重新选择")
    
    def voice_control_mode(self):
        """语音控制模式"""
        self.speak("请说话")
        audio_file = self.record_audio_direct(duration=5)
        
        if not audio_file:
            self.speak("没听到")
            return
        
        # 转录音频
        command = self.transcribe_audio(audio_file)
        
        if not command:
            self.speak("没听清")
            return
        
        print(f"🗣️  指令: {command}")
        self.speak(f"收到{command}")
        
        # 解析指令
        command_data = self.analyze_voice_command(command)
        
        # 显示解析结果
        action = command_data.get('action', '未知动作')
        duration = command_data.get('duration', 2)
        
        print(f"🔍 动作: {action}, 时长: {duration}秒")
        self.speak(f"准备{action}")
        
        # 快速确认
        self.speak("确认吗")
        confirm_audio = self.record_audio_direct(duration=3)
        if confirm_audio:
            confirm_text = self.transcribe_audio(confirm_audio)
            confirm = confirm_text if confirm_text else ""
        else:
            confirm = ""
        
        if confirm and any(word in confirm for word in ["确认", "执行", "是的", "好", "开始"]):
            self.execute_commands(command_data)
        else:
            self.speak("取消")
    
    def reset_arm(self):
        """复位机械臂到初始位置"""
        self.speak("复位中")
        reset_commands = {
            "action": "复位机械臂",
            "commands": [
                {"servo": "base", "position": 500},
                {"servo": "shoulder", "position": 500},
                {"servo": "elbow", "position": 500},
                {"servo": "wrist_vert", "position": 500},
                {"servo": "wrist_rot", "position": 500},
                {"servo": "gripper", "position": 0}
            ],
            "duration": 2
        }
        self.execute_commands(reset_commands)
    
    def read_servo_status(self):
        """读取舵机状态"""
        status_text = "状态: "
        for servo_name, position in self.arm_state.items():
            servo_id = self.servo_ids.get(servo_name, "未知")
            status_text += f"{servo_id}号{position}, "
        
        print(f"📊 {status_text}")
        self.speak(status_text)
    
    def single_servo_control(self):
        """单个舵机控制模式"""
        menu_options = {
            "1": "底座",
            "2": "肩部",
            "3": "肘部",
            "4": "腕部垂直", 
            "5": "腕部旋转",
            "6": "夹爪",
            "7": "返回"
        }
        
        choice = self.voice_menu_selection(menu_options, "选择舵机")
        
        servo_map = {
            "1": "base",
            "2": "shoulder",
            "3": "elbow", 
            "4": "wrist_vert",
            "5": "wrist_rot",
            "6": "gripper"
        }
        
        if choice in servo_map:
            servo_name = servo_map[choice]
            current_pos = self.arm_state[servo_name]
            
            self.speak(f"当前位置{current_pos}")
            pos_audio = self.record_audio_direct(duration=4)
            if pos_audio:
                pos_command = self.transcribe_audio(pos_audio)
            else:
                pos_command = ""
            
            if pos_command:
                numbers = re.findall(r'\d+', pos_command)
                if numbers:
                    pos = int(numbers[0])
                    if 0 <= pos <= 1000:
                        self.speak(f"设置{pos}")
                        self.control_single_servo(servo_name, pos)
                    else:
                        self.speak("范围错误")
                else:
                    self.speak("没听到数字")
        elif choice != "7" and choice != "cancel":
            self.speak("无效")

def main():
    controller = RoboticArmVoiceControl()
    controller.interactive_control()

if __name__ == "__main__":
    main()
